import Gnuplot
import paramiko
import argparse
import csv
import os
from time import sleep

initCalcsDone = False
numBodies = 0
fileLength = 0
chunkCount = 0
chunkStep = 0
currentChunk = 0
lastStep = -1
isDone = False
waitTime = 0.03
states = {}
plotList = []
autoscaled = 0

def loadNextChunk(f):
    global states
    global initCalcsDone
    global numBodies
    global fileLength
    global chunkCount
    global chunkStep
    global currentChunk
    global lastStep
    global isDone
    
    if not initCalcsDone:
        getInitVars(f)
        
    reader = csv.reader(f)
        
    if currentChunk < chunkCount:
        for numState in range(chunkStep):
            for b in range(numBodies):
                temp = reader.next()
                states[numState, b] = map(float, temp[1:4])
        currentChunk = currentChunk + 1
    elif currentChunk * chunkStep < fileLength:
        for numState in range(chunkStep):
            for b in range(numBodies):
                temp = reader.next()
                states[numState, b] = map(float, temp[1:4])
        lastStep = fileLength - (currentChunk * chunkStep)
    else:
        isDone = True

def getInitVars(f):
    global numBodies
    global fileLength
    global chunkCount
    global chunkStep
    global states
    global initCalcsDone

    reader = csv.reader(f)
    
    first = reader.next()
    
    if(len(first) == 2):
        numBodies = int(first[0])
        fileLength = int(first[1])
        chunkStep = numBodies * 10
        chunkCount = fileLength / chunkStep
    else:
        numBodies = int(first[0])
        fileLength = int(first[1])
        chunkCount = int(first[2])
        chunkStep = int(first[3])
        
    initCalcsDone = True

def animate(f,ranges):
    global states
    global lastStep
    global isDone
    global waitTime
    global plotList
    global autoscaled
    
    gplot = Gnuplot.Gnuplot()
    gplot.title('N-Body Simulation')
    gplot('set xrange [' + str(ranges[0]) + ' : ' + str(ranges[1]) + ']')
    gplot('set yrange [' + str(ranges[2]) + ' : ' + str(ranges[3]) + ']')
    gplot('set zrange [' + str(ranges[4]) + ' : ' + str(ranges[5]) + ']')
    gplot('set term x11 noraise')
    
    while not isDone:
        loadNextChunk(f)
        if isDone:
            break
        elif lastStep == -1:
            for i in range(chunkStep):
                fillPlotList(i)
                gplot.splot(plotList)
                plotList = []
                sleep(waitTime)
        elif lastStep > 0:
            for i in range(lastStep):
                fillPlotList(i)
                gplot.splot(plotList)
                plotList = []
                sleep(waitTime)
        else:
            break

def fillPlotList(i):
    global states
    global numBodies
    global plotList
    
    for j in range(numBodies):
        plotList.append(states[i, j])

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Animates a gravitational n-body simulation.')
    parser.add_argument('-f','--file',dest='localFile',type=argparse.FileType('rb'),help='File to use for the animation.')
    parser.add_argument('-s','--ssh',dest='confFile',default='',type=str,help='File that holds the ssh configuration.')
    parser.add_argument('-g','--get',dest='getPath',default='X',help='Filepath for the file on the ssh server.')
    parser.add_argument('-o','--output',dest='outPath',type=str,default='out.csv',help='Filepath to write sim data to.')
    parser.add_argument('-n',dest='useSFTP',action='store_false',default=True,help='Use to overide the use of SFTP (required when using -f)')
    parser.add_argument('--xmax',dest='xMaxRange',default=100,type=int,help='X max range')
    parser.add_argument('--xmin',dest='xMinRange',default=-100,type=int,help='X min range')
    parser.add_argument('--ymax',dest='yMaxRange',default=100,type=int,help='Y max range')
    parser.add_argument('--ymin',dest='yMinRange',default=-100,type=int,help='Y min range')
    parser.add_argument('--zmax',dest='zMaxRange',default=100,type=int,help='Z max range')
    parser.add_argument('--zmin',dest='zMinRange',default=-100,type=int,help='Z min range')
    parser.add_argument('--host',dest='host',help='Server address')
    parser.add_argument('--username',dest='username',help='Username for connecting to the server')
    parser.add_argument('--password',dest='password',help='Password for connecting to the server')
    parser.add_argument('--baseDir',dest='baseDir',type=str,help='Base directory for the server')
    parser.add_argument('-r','--ranges',dest='ranges',default=0,type=int,help='Set range scale')
    args = parser.parse_args()
    
    if not args.ranges == 0:
        args.xMaxRange = args.ranges
        args.xMinRange = -args.ranges
        args.yMaxRange = args.ranges
        args.yMinRange = -args.ranges
        args.zMaxRange = args.ranges
        args.zMinRange = -args.ranges
    
    r = [args.xMaxRange,args.xMinRange,args.yMaxRange,args.yMinRange,args.zMaxRange,args.zMinRange]
    
    if args.getPath == 'X' and args.useSFTP:
        print 'Error: need to specify getPath (-g, --get)'
        quit()
    
    if args.useSFTP:
        print('Connecting to remote host...')
        
        if not args.confFile == '':
            args.confFile = open(args.confFile,'rb')
            reader = csv.reader(args.confFile)
            conf = reader.next()
            
            host = conf[0]
            port = conf[1]
            username = conf[2]
            password = conf[3]
            baseDir = conf[4]
        else:
            host = args.host
            port = 22
            username = args.username
            password = args.password
            baseDir = args.baseDir
            
        try:
            transport = paramiko.Transport((host, int(port)))
            transport.connect(username = username, password = password)
            sftp = paramiko.SFTPClient.from_transport(transport)
            
            print 'Transfering data from the server... (takes longer for a bigger time amount)'
            sftp.get(args.getPath,args.outPath)
            
            sftp.close()
        except:
            print 'Error: cannot find server or ssh configuration is invalid!'
            quit()
            
        print 'Loading from file...'
        
        try:
            fout = open(args.outPath,'rb')
        except:
            print 'Error: failed to open file'
        animate(fout,r)
    else:
        print 'Loading from file...'
        animate(args.localFile,r)